// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <melon/benchmark.h>
#include <melon/init/init.h>
#include <pollux/connectors/hive/hive_partition_function.h>
#include <pollux/functions/lib/benchmarks/FunctionBenchmarkBase.h>
#include <pollux/type/type.h>
#include <pollux/vector/complex_vector.h>
#include <pollux/vector/vector_fuzzer.h>
#include <pollux/vector/vector_maker.h>

TURBO_FLAG(int64_t,fuzzer_seed, 99887766, "Seed for random input dataset generator");

using namespace kumo::pollux;
using namespace kumo::pollux::test;
using connector::hive::HivePartitionFunction;

namespace {

constexpr std::array<TypeKind, 10> kSupportedScalarTypes{
    TypeKind::BOOLEAN,
    TypeKind::TINYINT,
    TypeKind::SMALLINT,
    TypeKind::INTEGER,
    TypeKind::BIGINT,
    TypeKind::REAL,
    TypeKind::DOUBLE,
    TypeKind::VARCHAR,
    TypeKind::TIMESTAMP};

class HivePartitionFunctionBenchmark
    : public functions::test::FunctionBenchmarkBase {
 public:
  explicit HivePartitionFunctionBenchmark(size_t vectorSize)
      : FunctionBenchmarkBase() {
    // Prepare input data
    VectorFuzzer::Options opts;
    opts.vectorSize = vectorSize;
    opts.stringLength = 20;
    VectorFuzzer fuzzer(opts, pool(), FLAGS_fuzzer_seed);
    VectorMaker vm{pool_.get()};
    auto addRowVector = [&](const TypePtr& type) {
      auto flat_vector = fuzzer.fuzzFlat(type);
      rowVectors_[type->kind()] = vm.row_vector({flat_vector});
    };

    for (auto typeKind : kSupportedScalarTypes) {
      addRowVector(createScalarType(typeKind));
    }

    addRowVector(ARRAY(REAL()));
    addRowVector(MAP(BIGINT(), BOOLEAN()));
    addRowVector(ROW({"a", "b"}, {INTEGER(), DOUBLE()}));

    // Prepare HivePartitionFunction
    fewBucketsFunction_ = createHivePartitionFunction(20);
    manyBucketsFunction_ = createHivePartitionFunction(100);

    partitions_.resize(vectorSize);
  }

  template <TypeKind KIND>
  void runFew() {
    run<KIND>(fewBucketsFunction_.get());
  }

  template <TypeKind KIND>
  void runMany() {
    run<KIND>(manyBucketsFunction_.get());
  }

 private:
  std::unique_ptr<HivePartitionFunction> createHivePartitionFunction(
      size_t bucketCount) {
    std::vector<int> bucketToPartition(bucketCount);
    std::iota(bucketToPartition.begin(), bucketToPartition.end(), 0);
    std::vector<column_index_t> keyChannels;
    keyChannels.emplace_back(0);
    return std::make_unique<HivePartitionFunction>(
        bucketCount, bucketToPartition, keyChannels);
  }

  template <TypeKind KIND>
  void run(HivePartitionFunction* function) {
    if (rowVectors_.find(KIND) == rowVectors_.end()) {
      throw std::runtime_error(
          fmt::format("Unsupported type {}.", mapTypeKindToName(KIND)));
    }
    function->partition(*rowVectors_[KIND], partitions_);
  }

  std::unordered_map<TypeKind, RowVectorPtr> rowVectors_;
  std::unique_ptr<HivePartitionFunction> fewBucketsFunction_;
  std::unique_ptr<HivePartitionFunction> manyBucketsFunction_;
  std::vector<uint32_t> partitions_;
};

std::unique_ptr<HivePartitionFunctionBenchmark> benchmarkFew;
std::unique_ptr<HivePartitionFunctionBenchmark> benchmarkMany;

BENCHMARK(booleanFewRowsFewBuckets) {
  benchmarkFew->runFew<TypeKind::BOOLEAN>();
}

BENCHMARK_RELATIVE(booleanFewRowsManyBuckets) {
  benchmarkFew->runMany<TypeKind::BOOLEAN>();
}

BENCHMARK(booleanManyRowsFewBuckets) {
  benchmarkMany->runFew<TypeKind::BOOLEAN>();
}

BENCHMARK_RELATIVE(booleanManyRowsManyBuckets) {
  benchmarkMany->runMany<TypeKind::BOOLEAN>();
}

BENCHMARK_DRAW_LINE();

BENCHMARK(tinyintFewRowsFewBuckets) {
  benchmarkFew->runFew<TypeKind::TINYINT>();
}

BENCHMARK_RELATIVE(tinyintFewRowsManyBuckets) {
  benchmarkFew->runMany<TypeKind::TINYINT>();
}

BENCHMARK(tinyintManyRowsFewBuckets) {
  benchmarkMany->runFew<TypeKind::TINYINT>();
}

BENCHMARK_RELATIVE(tinyintManyRowsManyBuckets) {
  benchmarkMany->runMany<TypeKind::TINYINT>();
}

BENCHMARK_DRAW_LINE();

BENCHMARK(smallintFewRowsFewBuckets) {
  benchmarkFew->runFew<TypeKind::SMALLINT>();
}

BENCHMARK_RELATIVE(smallintFewRowsManyBuckets) {
  benchmarkFew->runMany<TypeKind::SMALLINT>();
}

BENCHMARK(smallintManyRowsFewBuckets) {
  benchmarkMany->runFew<TypeKind::SMALLINT>();
}

BENCHMARK_RELATIVE(smallintManyRowsManyBuckets) {
  benchmarkMany->runMany<TypeKind::SMALLINT>();
}

BENCHMARK_DRAW_LINE();

BENCHMARK(integerFewRowsFewBuckets) {
  benchmarkFew->runFew<TypeKind::INTEGER>();
}

BENCHMARK_RELATIVE(integerFewRowsManyBuckets) {
  benchmarkFew->runMany<TypeKind::INTEGER>();
}

BENCHMARK(integerManyRowsFewBuckets) {
  benchmarkMany->runFew<TypeKind::INTEGER>();
}

BENCHMARK_RELATIVE(integerManyRowsManyBuckets) {
  benchmarkMany->runMany<TypeKind::INTEGER>();
}

BENCHMARK_DRAW_LINE();

BENCHMARK(bigintFewRowsFewBuckets) {
  benchmarkFew->runFew<TypeKind::BIGINT>();
}

BENCHMARK_RELATIVE(bigintFewRowsManyBuckets) {
  benchmarkFew->runMany<TypeKind::BIGINT>();
}

BENCHMARK(bigintManyRowsFewBuckets) {
  benchmarkMany->runFew<TypeKind::BIGINT>();
}

BENCHMARK_RELATIVE(bigintManyRowsManyBuckets) {
  benchmarkMany->runMany<TypeKind::BIGINT>();
}

BENCHMARK_DRAW_LINE();

BENCHMARK(realFewRowsFewBuckets) {
  benchmarkFew->runFew<TypeKind::REAL>();
}

BENCHMARK_RELATIVE(realFewRowsManyBuckets) {
  benchmarkFew->runMany<TypeKind::REAL>();
}

BENCHMARK(realManyRowsFewBuckets) {
  benchmarkMany->runFew<TypeKind::REAL>();
}

BENCHMARK_RELATIVE(realManyRowsManyBuckets) {
  benchmarkMany->runMany<TypeKind::REAL>();
}

BENCHMARK_DRAW_LINE();

BENCHMARK(doubleFewRowsFewBuckets) {
  benchmarkFew->runFew<TypeKind::DOUBLE>();
}

BENCHMARK_RELATIVE(doubleFewRowsManyBuckets) {
  benchmarkFew->runMany<TypeKind::DOUBLE>();
}

BENCHMARK(doubleManyRowsFewBuckets) {
  benchmarkMany->runFew<TypeKind::DOUBLE>();
}

BENCHMARK_RELATIVE(doubleManyRowsManyBuckets) {
  benchmarkMany->runMany<TypeKind::DOUBLE>();
}

BENCHMARK_DRAW_LINE();

BENCHMARK(varcharFewRowsFewBuckets) {
  benchmarkFew->runFew<TypeKind::VARCHAR>();
}

BENCHMARK_RELATIVE(varcharFewRowsManyBuckets) {
  benchmarkFew->runMany<TypeKind::VARCHAR>();
}

BENCHMARK(varcharManyRowsFewBuckets) {
  benchmarkMany->runFew<TypeKind::VARCHAR>();
}

BENCHMARK_RELATIVE(varcharManyRowsManyBuckets) {
  benchmarkMany->runMany<TypeKind::VARCHAR>();
}

BENCHMARK_DRAW_LINE();

BENCHMARK(timestampFewRowsFewBuckets) {
  benchmarkFew->runFew<TypeKind::TIMESTAMP>();
}

BENCHMARK_RELATIVE(timestampFewRowsManyBuckets) {
  benchmarkFew->runMany<TypeKind::TIMESTAMP>();
}

BENCHMARK(timestampManyRowsFewBuckets) {
  benchmarkMany->runFew<TypeKind::TIMESTAMP>();
}

BENCHMARK_RELATIVE(timestampManyRowsManyBuckets) {
  benchmarkMany->runMany<TypeKind::TIMESTAMP>();
}

BENCHMARK_DRAW_LINE();

BENCHMARK(arrayFewRowsFewBuckets) {
  benchmarkFew->runFew<TypeKind::ARRAY>();
}

BENCHMARK_RELATIVE(arrayFewRowsManyBuckets) {
  benchmarkFew->runMany<TypeKind::ARRAY>();
}

BENCHMARK(arrayManyRowsFewBuckets) {
  benchmarkMany->runFew<TypeKind::ARRAY>();
}

BENCHMARK_RELATIVE(arrayManyRowsManyBuckets) {
  benchmarkMany->runMany<TypeKind::ARRAY>();
}

BENCHMARK_DRAW_LINE();

BENCHMARK(mapFewRowsFewBuckets) {
  benchmarkFew->runFew<TypeKind::MAP>();
}

BENCHMARK_RELATIVE(mapFewRowsManyBuckets) {
  benchmarkFew->runMany<TypeKind::MAP>();
}

BENCHMARK(mapManyRowsFewBuckets) {
  benchmarkMany->runFew<TypeKind::MAP>();
}

BENCHMARK_RELATIVE(mapManyRowsManyBuckets) {
  benchmarkMany->runMany<TypeKind::MAP>();
}

BENCHMARK_DRAW_LINE();

BENCHMARK(rowFewRowsFewBuckets) {
  benchmarkFew->runFew<TypeKind::ROW>();
}

BENCHMARK_RELATIVE(rowFewRowsManyBuckets) {
  benchmarkFew->runMany<TypeKind::ROW>();
}

BENCHMARK(rowManyRowsFewBuckets) {
  benchmarkMany->runFew<TypeKind::ROW>();
}

BENCHMARK_RELATIVE(rowManyRowsManyBuckets) {
  benchmarkMany->runMany<TypeKind::ROW>();
}

BENCHMARK_DRAW_LINE();
} // namespace

int main(int argc, char** argv) {
  melon::Init init{&argc, &argv};
  gflags::ParseCommandLineFlags(&argc, &argv, true);
  memory::MemoryManager::initialize({});
  benchmarkFew = std::make_unique<HivePartitionFunctionBenchmark>(1'000);
  benchmarkMany = std::make_unique<HivePartitionFunctionBenchmark>(10'000);

  melon::runBenchmarks();

  benchmarkFew.reset();
  benchmarkMany.reset();

  return 0;
}
